#
# Benchmark where a message (an integer) is repeatedly passed aroung a
# ring of processes/threads.
#
# Each time the message arrives at the process/thread at position 0
# the message gets decremented.  Whenever a process/thread get the
# message of value 0 it passes it on to its successor and then exits.
#
# Compare with `examples/ring_c.c` in the OpenMPI distribution.
#
# done by Richard Oudkerk

import processing
import threading
import Queue
import time

#
# Function run by processes/threads in ring
#

def f(i, prev, next):
    if hasattr(prev, 'get'):
        get, put = prev.get, next.put
    else:
        get, put = prev.recv, next.send

    msg = True
    while msg:
        msg = get()
        if i == 0:
            msg -= 1
        put(msg)

#
# Functions to create a ring and associated start function
#

def PipeRing(size, times):
    pipes = [processing.Pipe() for i in range(size)]
    ring = [processing.Process(target=f, args=(i, pipes[i-1][1], pipes[i][0]))
            for i in range(size)]
    def starter():
        return pipes[-1][0].send(times)
    return ring, starter

def QueueRing(size, times, SpawnType, QueueType):
    queues = [QueueType() for i in range(size)]
    ring = [SpawnType(target=f, args=(i, queues[i-1], queues[i]))
            for i in range(size)]
    def starter():
        return queues[-1].put(times)
    return ring, starter

#
# Parameters
#

RINGSIZE = 8    # number of threads/processes in ring
TIMES = 10000   # the number of times message will pass around ring
                #     (this is also the initial value of the message)

params = [
    ('processing.Process and processing.Pipe',
         PipeRing, []),
    ('processing.Process and processing.PipeQueue',
         QueueRing, [processing.Process, processing.PipeQueue]),
    ('threading.Threads and Queue.Queue',
         QueueRing, [threading.Thread, Queue.Queue])         
    ]

if hasattr(processing, 'PosixQueue'):
    params.insert(0, 
        ('processing.Process and processing.PosixQueue',
             QueueRing, [processing.Process, processing.PosixQueue]))

#
# Main
#

if __name__ == '__main__':
    print 'RINGSIZE = %d, TIMES = %d\n' % (RINGSIZE, TIMES)

    for descr, Ring, args in params:
        ring, start_passing = Ring(RINGSIZE, TIMES, *args)

        for r in ring:           # start the threads/processes
            r.start()

        time.sleep(1)            # wait for threads/processes to start up

        t = time.time()
        
        start_passing()          # give intial message to ring

        for r in ring:           # wait for the threads/processes to finish
            r.join()
        
        elapsed = time.time() - t
        
        print descr
        print '\ttime spent: %g' % elapsed
        print '\tmessages/sec: %g' % ((RINGSIZE*TIMES)/elapsed)
        print
